import pandas as pd
import torch
import torch.nn as nn
from torch.optim import AdamW
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModel

"""读取数据，查看数据情况"""
data_path = 'data/fake_or_real_news.csv'
data = pd.read_csv(data_path)  # 新闻 data['id','title','text','label']
# 获取新闻的标题、文本和标签（真假）,
titles = data['title'].tolist()
texts = data['text'].tolist()
labels = data['label'].tolist()
print(f'样本总数为：{len(titles)}')

"""划分数据集、测试集"""
"""
from sklearn.model_selection import StratifiedKFold
skf = StratifiedKFold(n_splits=10, random_state=123, shuffle=True)
n_splits: 这是要生成的交叉验证折数，即数据集将被分成的子集数量。在这里，指定为 10，表示数据集将被分成 10 折进行交叉验证。
random_state: 这个参数是用来控制随机数生成器行为的种子，以确保每次运行时分割的结果都是相同的。
shuffle: 这是一个布尔值参数，指示是否在进行交叉验证之前对数据进行洗牌。

# 假设有一个特征数组 X 和一个目标数组 y
X = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]
y = [0, 0, 1, 1, 1]

# 创建 Stratified K-Folds 对象
skf = StratifiedKFold(n_splits=10, random_state=123, shuffle=True)

# 使用 split 方法来获取每一折的训练集和测试集索引
for train_index, test_index in skf.split(X, y):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]
    print("Train Index:", train_index)
    print("Test Index:", test_index)
Train Index: [0 1 3 4]
Test Index: [2]
Train Index: [1 2 3 4]
Test Index: [0]
Train Index: [0 2 3]
Test Index: [1 4]
    # 在这里可以对 X_train, y_train 进行训练，并对 X_test, y_test 进行测试
"""
from sklearn.model_selection import StratifiedKFold

skf = StratifiedKFold(n_splits=3)
for train_ids, test_ids in skf.split(texts, labels):
    # train_ids 训练集的新闻索引
    # test_ids  测试集的新闻索引
    train_X = [texts[i] for i in train_ids]
    test_X = [texts[i] for i in test_ids]
    train_Y = [labels[i] for i in train_ids]
    test_Y = [labels[i] for i in test_ids]
    break
print(len(train_X), len(train_X), len(test_Y), len(test_Y))

"""数据数值化。覆写torch下的Dataset类，用于得到模型的输入"""

from torch.utils.data import Dataset


class MyDataset(Dataset):
    def __init__(self, X, Y):
        self.tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
        self.X = X
        self.Y = Y

    def __len__(self):
        return len(self.Y)

    def __getitem__(self, idx):
        text = self.X[idx]
        label = self.Y[idx]
        """
        tokenizer 的作用是将文本转换为 BERT 模型可以接受的输入格式，包括将文本分词并转换成对应的 token IDs，
        以及生成 attention mask 等功能。
        
        文本张量化
        """
        inputs = self.tokenizer(text, return_tensors='pt', max_length=200, padding="max_length", truncation=True)
        target = ['FAKE', 'REAL'].index(label)  # 将标签变为1 0
        # inputs = {"input_ids":tensor,"token_type_ids":tensor,"attention_mask":tensor}
        return {"inputs": inputs['input_ids'].squeeze(), "targets": target}

    """
    作用：
    "[cls(标志这句话)]我是中国人"————>[101(cls的标记),232,345,554,656,223]  文本张量化
    "FAKE"——>0
    "REAL"——>1
    """


# 加载数据
from torch.utils.data import DataLoader

my_dataset = MyDataset(train_X, train_Y)
train_dataloader = DataLoader(my_dataset, batch_size=8, shuffle=True)
test_dataloader = DataLoader(my_dataset, batch_size=8, shuffle=True)
for batch in train_dataloader:
    print(batch['inputs'].shape)
    print(batch['targets'].shape)
    break

"""
训练
1.设备定义
2.模型
3.损失函数
4.优化器
5.训练数据
"""


# 构建模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.plm = AutoModel.from_pretrained('bert-base-uncased')  # 用于加载预训练的 BERT 模型
        # from_pretrained() 是 AutoModel 类的一个方法，用于从 Hugging Face 模型存储库中加载预训练模型。
        # 'bert-base-uncased' 指定了一个预训练的 BERT 模型，该模型是基于小写字母的版本。这个模型是在大规模文本语料上进行预训练的，包括了词汇编码和语言理解的能力。
        self.cls = nn.Linear(768, 2)
        """
        nn.Linear 是 PyTorch 中用于定义线性层的类。
        构造函数的参数是 (in_features, out_features)，分别指定输入特征的维度和输出特征的维度。
        768 是输入特征的维度，通常对应于 BERT 模型的隐藏状态的维度。
        2 是输出特征的维度，这里看起来是指定输出的类别数，因此这是一个用于分类任务的线性层。
        
        在自然语言处理中，特别是在基于预训练模型（如 BERT）进行下游任务（如文本分类）时，
        通常会使用一个预训练的模型来提取文本特征，然后通过线性层来进行分类。
        在这种情况下，BERT 模型通常输出一个隐藏状态向量，其维度为 768（或其他具体模型对应的维度），
        然后通过 nn.Linear(768, 2) 将这个向量映射到一个维度为 2 的输出，用来表示两个类别的分类概率。
        """

    def forward(self, inputs):  # forward 方法接受一个输入 inputs（假设是 BERT 模型可以接受的输入序列），通过 self.plm 提取特征并通过 self.cls 进行分类预测。
        outputs = self.plm(inputs).pooler_output
        outputs = self.cls(outputs)
        return outputs
        # outputs.shape = (8,2)  返回的是8个样本每个样本[FAKE,REAL] 对应的概率


import torch.nn.functional as F


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 使用HuggingFace的transformers库中的AutoModel加载预训练的BERT模型
        # 使用'bert-base-uncased'预训练模型，这是一个小写版本的BERT模型。
        self.plm = AutoModel.from_pretrained('bert-base-uncased')  # 文本——>对应向量
        # 将BERT模型的所有参数设为不可训练，即冻结BERT模型的参数，以免在训练过程中更新这些参数。
        for param in self.plm.parameters():
            param.requires_grad = False  # 预训练的BERT模型的参数不参与梯度计算
        # 定义了一个 ModuleList，里面包含了三个二维卷积层 (Conv2d)。
        # 每个卷积层的输入通道是1，输出通道是64，卷积核大小分别为 (4, 768), (8, 768), 和 (15, 768)。
        """
        [4,768]的卷积核，与[7681]的文本序列卷积后得到[4,1]的结果
        则[64,4,768]的卷积（三维图形）,与[7681]的文本序列卷积后得到[64,4,1]的结果————输出为64个[4,1]的结果
        """
        self.convs = nn.ModuleList(
            [nn.Conv2d(1, 64, (4, 768)),  # 768为文本序列长度，包含所有列，为768
             nn.Conv2d(1, 64, (8, 768)),
             nn.Conv2d(1, 64, (15, 768)), ]
        )
        # 定义了一个线性层 (Linear)，输入特征数是64 * 3，输出特征数是2，用于最终的分类任务。
        self.cls = nn.Linear(64 * 3, 2)

    def conv_and_pool(self, x, conv):  # x.shape = [8,200,768] 与 Conv2d(1, 64, (4, 768)卷积
        # 对输入 x 进行卷积操作 conv(x)，得到[8,64,200-4+1,1]  8个样本，每个样本有64个卷积得到的向量组成的矩阵
        # 然后通过 squeeze(3) 去掉维度为3的尺寸，即将宽度为1的维度压缩掉，然后使用ReLU激活函数。
        x = F.relu(conv(x).squeeze(3))  # x.shape = [8,64,200-4+1]
        # 对经过ReLU激活的卷积结果进行一维最大池化操作 (max_pool1d)，池化的窗口大小为 x.size(2)【对64个得到卷积结果每一个进行池化操作得到[64,1]】，
        # 然后通过 squeeze(2) 去掉维度为2的尺寸，即将高度为1的维度压缩掉，得到池化后的结果 x
        x = F.max_pool1d(x, x.size(2)).squeeze(2)  # x.shape = [8,64]
        return x

    def forward(self, batch_inputs):  # batch_inputs.shape=[8,200] 8个样本，每个样本为200个token
        # 将输入的文本ID batch_inputs 输入到BERT模型 self.plm 中，
        # 获取BERT模型最后一层的隐藏状态 last_hidden_state
        encoder_out = self.plm(input_ids=batch_inputs).last_hidden_state # [8,200,768] 8个样本，每个样本为200个token，每个token被嵌入768个向量
        # 在 encoder_out 张量的第一维度上增加一个维度，以便适配二维卷积的输入格式。
        out = encoder_out.unsqueeze(1)  # [8,1,200,768]
        # 将out[8,1,200,768]与三个卷积层分别进行卷积和池化操作，并使用 torch.cat 沿着第二维度（列方向）拼接起来，形成一个特征向量
        # 每一个卷积层与out卷积和池化操作得到[8,64]向量，三个拼接为[8,64*3]
        out = torch.cat([self.conv_and_pool(out, conv) for conv in self.convs])
        # 将拼接后的特征向量 out 输入到线性层 self.cls 进行分类预测，得到最终的输出结果 out
        # [8,64*3]的结果与[64*3,2]的全连接层进行计算得到[8,2]8个样本的预测结果
        out = self.cls(out)
        return out


class LSTM(nn.Module):
    def __init__(self):
        super(LSTM, self).__init__()
        self.plm = AutoModel.from_pretrained('bert-base-uncased')  # 文本——>对应向量
        for param in self.plm.parameters():
            param.requires_grad = False  # 预训练的BERT模型的参数不参与梯度计算
        # 定义了一个双向 LSTM 模型，输入维度为 768（BERT 的输出维度），隐藏层维度为 200
        # batch_first=True 表示输入数据的第一个维度是 batch_size。
        self.lstm = nn.LSTM(768, 200, 1, batch_first=True, bidirectional=True)
        # 最终输出一个隐藏层向量作为输出层，维度为200      LTSM采用前后时间点的隐藏层，即维度为200*2
        # 定义了一个线性层，输入大小为 200 * 2，输出大小为 2。这个线性层用于最终的分类任务。
        self.fc = nn.Linear(200 * 2, 2)

    def forward(self, batch_inputs):
        # 将输入的 batch_inputs 送入 BERT 模型，得到 BERT 的最后一个隐藏状态作为编码器的输出。
        encoder_out = self.plm(input_ids=batch_inputs).last_hidden_state
        # 将 BERT 的输出作为 LSTM 的输入，得到 LSTM 的输出 out。_ 是 LSTM 的隐藏状态，但在这个例子中没有使用。
        out, _ = self.lstm(encoder_out)
        # 将 LSTM 的输出的最后一个时间步的输出（即 out[:, -1, :]）通过线性层 self.fc 进行分类预测，得到最终的输出 out
        out = self.fc(out[:, -1, :])
        return out


"""
self.plm(input_ids=inputs):

这里 self.plm 是一个预训练的BERT语言模型
input_ids 是语言模型接受的输入，通常是 tokenized 后的文本数据的索引序列。
self.plm(input_ids=inputs) 的调用将输入 inputs传递给 self.plm，并期望获得模型的输出。

.pooler_output:
假设 self.plm 的输出中有一个 pooler_output 属性，它是经过池化处理后的模型输出。
在 BERT 模型中，pooler_output 是一个经过特殊处理的表示[cls]，通常被用作整个句子的表示或者上下文信息。

self.cls(outputs):
outputs 是经过语言模型处理后得到的特征表示，即 pooler_output。
self.cls 是一个定义好的线性层，将输入的特征维度映射到输出的维度（可能是用于分类的维度）。
"""
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# 使用模型
model = MyModel().to(device)
# 损失函数
loss_fn = nn.CrossEntropyLoss().to(device)  # 分类问题一般使用交叉熵损失函数
# 定义优化器
optimizer = AdamW(model.parameters(), lr=0.0001)


# 定义训练函数
def train(device, model, loss_fn, optimizer, train_dataloader, epoch):
    model.train()
    # 对model进行训练，计算损失
    """
    1.依次拿出一个batch的样本
    2.将该样本的inputs输入到模型中
    3.计算当前batch的模型输出
    4.用损失函数计算模型输出与真实标签之间的损失
    5.损失对权重反向传播求梯度
    6.更新权重
    7.计算当前的平均准确率和当前平均损失
    """
    num_examples = 0  # 模型样本总数
    total_correct = 0  # 总共预测正确个数
    total_loss = 0  # 模型样本总损失
    bar = tqdm(enumerate(train_dataloader))  # 增加进度条
    """
    enumerate(train_dataloader)：
    train_dataloader 是一个 PyTorch 的 DataLoader 对象，用于加载训练数据批次。
    enumerate(train_dataloader) 将 train_dataloader 转换为一个可枚举的迭代器，每次迭代返回一个元组 (batch_idx, batch_data)，
    其中 batch_idx 是批次的索引，batch_data 是对应的数据批次。
    
    tqdm() 函数：
    tqdm 是一个 Python 进度条库，用于在循环或迭代过程中显示进度条。
    当你调用 tqdm(enumerate(train_dataloader))，它会接收一个可枚举的对象，并返回一个带有进度条的迭代器对象 bar。
    """
    # 枚举返回 标号和train_dataloader
    for i, batch in bar:
        batch_inputs = batch['inputs'].to(device)
        batch_labels = batch['targets'].to(device)
        # 获取模型并计算损失
        out = model(batch_inputs)
        loss = loss_fn(out, batch_labels)
        # 优化  1.清空梯度  2.反向传播  3.优化参数
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 计算准确率和平均损失
        num_examples += len(batch_labels)  # 目前模型训练的样本总数
        batch_pred = out.argmax(dim=-1)  # 找到概率大的标签作为预测值[8,1]
        correct = (batch_pred == batch_labels).sum().item()  # 此批次预测正确的个数
        total_correct += correct  # 所有预测正确的个数
        accuracy = total_correct / num_examples  # 预测正确率
        total_loss += loss.item()  # 目前样本预测的损失和
        avg_loss = total_loss / num_examples  # 平均损失

        # set_postfix() 方法允许你在每次迭代过程中动态更新显示在进度条后缀的信息。
        # 通常使用键值对的方式传递信息，方便 tqdm 更新后缀信息{epoch:epoch,accuracy:accuracy,avg_loss:avg_loss}
        bar.set_postfix(epoch=epoch, accuracy=accuracy, avg_loss=avg_loss)

    return accuracy, avg_loss


# 定义测试函数
def test(device, model, loss_fn, test_dataloader, epoch):
    model.eval()
    # 对model进行训练，计算损失
    """
    1.依次拿出一个batch的样本
    2.将该样本的inputs输入到模型中
    3.计算当前batch的模型输出
    4.用损失函数计算模型输出与真实标签之间的损失
    5.计算当前的平均准确率和当前平均损失
    """
    num_examples = 0  # 模型样本总数
    total_correct = 0  # 总共预测正确个数
    total_loss = 0  # 模型样本总损失
    bar = tqdm(enumerate(test_dataloader))  # 增加进度条
    # 枚举返回 标号和train_dataloader
    for i, batch in bar:
        batch_inputs = batch['inputs'].to(device)
        batch_labels = batch['targets'].to(device)
        # 获取模型并计算损失
        out = model(batch_inputs)
        loss = loss_fn(out, batch_labels)

        # 计算准确率和平均损失
        num_examples += len(batch_labels)  # 目前模型训练的样本总数
        batch_preds = out.argmax(dim=-1)  # 找到概率大的标签作为预测值[8,1]
        correct = (batch_preds == batch_labels).sum().item()  # 此批次预测正确的个数
        total_correct += correct  # 所有预测正确的个数
        accuracy = total_correct / num_examples  # 预测正确率
        total_loss += loss.item()  # 目前样本预测的损失和
        avg_loss = total_loss / num_examples  # 平均损失

        # set_postfix() 方法允许你在每次迭代过程中动态更新显示在进度条后缀的信息。
        # 通常使用键值对的方式传递信息，方便 tqdm 更新后缀信息{epoch:epoch,accuracy:accuracy,avg_loss:avg_loss}
        bar.set_postfix(epoch=epoch, accuracy=accuracy, avg_loss=avg_loss)

    return accuracy, avg_loss


# 开始训练
epochs = 10  # 训练轮次

for epoch in range(1, epochs + 1):
    train_accuracy, train_loss = train(device, model, loss_fn, optimizer, train_dataloader, epoch)
    test_accuracy, test_loss = test(device, model, loss_fn, train_dataloader, epoch)
